import pandas as pd
import os
from datetime import datetime, timedelta
import plotly.graph_objects as go
from plotly.subplots import make_subplots


# --- 数据加载和处理函数 (与 draw.py 相同) ---
def load_data_in_range(start_time_str: str, end_time_str: str, base_path: str, symbol: str,
                       timeframe: str) -> pd.DataFrame:
    """从指定的CSV文件中加载时间范围内的数据，并动态计算Volume Delta。"""
    start_date = pd.to_datetime(start_time_str)
    end_date = pd.to_datetime(end_time_str)
    data_folder_path = os.path.join(base_path, symbol, timeframe)
    if not os.path.isdir(data_folder_path):
        print(f"错误: 文件夹不存在 -> {data_folder_path}")
        return pd.DataFrame()

    date_range = pd.date_range(start=start_date.date(), end=end_date.date(), freq='D').strftime('%Y-%m').unique()
    if len(date_range) == 0: date_range = [start_date.strftime('%Y-%m')]
    all_dfs = []
    print(f"将在以下月份文件中查找数据: {list(date_range)}")
    for year_month in date_range:
        file_path = os.path.join(data_folder_path, f"{symbol.upper()}-{timeframe}-{year_month}.csv")
        if os.path.exists(file_path):
            print(f"正在读取文件: {file_path}")
            try:
                df = pd.read_csv(file_path)
                required_cols = ['open_time', 'open', 'high', 'low', 'close', 'volume', 'taker_buy_volume']
                if not all(col in df.columns for col in required_cols):
                    print(f"错误: 文件 {file_path} 缺少必要的列。")
                    continue

                for col in ['volume', 'taker_buy_volume']:
                    df[col] = pd.to_numeric(df[col], errors='coerce')

                df['volume_delta'] = (2 * df['taker_buy_volume']) - df['volume']
                all_dfs.append(df)
            except Exception as e:
                print(f"读取或处理文件 {file_path} 时出错: {e}")
        else:
            print(f"警告: 未找到文件 {file_path}")

    if not all_dfs:
        print("错误: 在指定时间范围内未找到任何数据文件。")
        return pd.DataFrame()

    combined_df = pd.concat(all_dfs, ignore_index=True)
    combined_df.drop_duplicates(subset=['open_time'], inplace=True)
    return combined_df


def process_and_filter_data(df: pd.DataFrame, start_time_str: str, end_time_str: str, 
                            include_warmup: bool = False, warmup_periods: int = 0) -> pd.DataFrame:
    """
    处理并筛选数据。
    
    参数:
        df: 原始数据 DataFrame
        start_time_str: 开始时间字符串
        end_time_str: 结束时间字符串
        include_warmup: 是否包含预热期数据
        warmup_periods: 预热期K线数量
    """
    if df.empty:
        return df
    df['open_time'] = pd.to_datetime(df['open_time'], unit='ms')
    df.set_index('open_time', inplace=True)
    numeric_cols = ['open', 'high', 'low', 'close', 'volume', 'volume_delta']
    df[numeric_cols] = df[numeric_cols].apply(pd.to_numeric, errors='coerce')
    df.dropna(subset=numeric_cols, inplace=True)
    df.sort_index(inplace=True)
    df['cvd'] = df['volume_delta'].cumsum()
    
    # 使用 pd.to_datetime 转换字符串时间以确保时区一致性
    start_ts = pd.to_datetime(start_time_str)
    end_ts = pd.to_datetime(end_time_str)
    
    if include_warmup and warmup_periods > 0:
        # 如果需要包含预热期，找到开始时间之前（严格小于）的 warmup_periods 根K线
        data_before_start = df[df.index < start_ts]
        if len(data_before_start) >= warmup_periods:
            # 取最后 warmup_periods 根K线作为预热期
            warmup_start_idx = data_before_start.index[-warmup_periods]
            df_filtered = df.loc[warmup_start_idx:end_ts]
            actual_warmup = len(df_filtered[df_filtered.index < start_ts])
            print(f"数据处理完成。包含 {actual_warmup} 根预热期K线，筛选出 {len(df_filtered)} 条K线（含预热期）。")
        else:
            # 如果预热期数据不足，使用所有可用数据
            df_filtered = df.loc[:end_ts]
            actual_warmup = len(data_before_start)
            print(f"数据处理完成。可用预热期数据: {actual_warmup} 根（请求: {warmup_periods}），筛选出 {len(df_filtered)} 条K线。")
    else:
        df_filtered = df.loc[start_ts:end_ts]
        print(f"数据处理完成。筛选出 {len(df_filtered)} 条K线。")
    
    return df_filtered


# --- Kyle's Lambda 计算函数 ---
def calculate_kyles_lambda(df: pd.DataFrame, window: int = 20) -> pd.Series:
    """
    计算 Kyle's Lambda 指标。
    
    Kyle's Lambda 衡量价格对订单流的敏感度，通过线性回归计算：
    Lambda = Cov(Delta P, Net Flow) / Var(Net Flow)
    
    参数:
        df: 包含 'close', 'open', 'volume_delta' 列的 DataFrame
        window: 滚动窗口大小（默认20根K线）
    
    返回:
        包含 Kyle's Lambda 值的 Series
    """
    if df.empty:
        return pd.Series(dtype=float)
    
    # 1. 计算价格变动 (Delta P)
    df = df.copy()  # 避免修改原始 DataFrame
    df['delta_p'] = df['close'] - df['open']
    
    # 2. 净订单流 (Net Order Flow Q) 已经计算为 volume_delta
    # net_flow = volume_delta = 2 * taker_buy_volume - volume
    
    # 3. 使用滚动窗口计算 Kyle's Lambda
    kyles_lambda = []
    
    for i in range(len(df)):
        if i < window - 1:
            # 窗口数据不足，设为 NaN
            kyles_lambda.append(pd.NA)
            continue
        
        # 获取窗口数据
        window_data = df.iloc[i - window + 1:i + 1]
        
        # 检查数据有效性
        variance = window_data['volume_delta'].var()
        if pd.isna(variance) or variance == 0:
            kyles_lambda.append(0.0)
            continue
        
        # 计算协方差和方差
        covariance = window_data['delta_p'].cov(window_data['volume_delta'])
        
        if pd.isna(covariance) or variance == 0:
            kyles_lambda.append(0.0)
        else:
            kyles_lambda.append(covariance / variance)
    
    return pd.Series(kyles_lambda, index=df.index)


# --- 绘图函数 (K线图 + Kyle's Lambda 合并) ---
def plot_with_kyles_lambda(df: pd.DataFrame, chart_title: str, lambda_window: int = 20, 
                           start_time_str: str = None, warmup_periods: int = 0):
    """
    绘制K线图，并将 Kyle's Lambda 合并到同一图表中（使用双Y轴）。
    
    参数:
        df: 包含K线数据的 DataFrame（可能包含预热期数据）
        chart_title: 图表标题
        lambda_window: Kyle's Lambda 的滚动窗口大小
        start_time_str: 实际开始时间（用于区分预热期和正式数据）
        warmup_periods: 预热期K线数量
    """
    if df.empty:
        print("数据为空，无法绘制图表。")
        return
    
    # 计算 Kyle's Lambda
    print(f"\n--- 正在计算 Kyle's Lambda (窗口大小: {lambda_window}) ---")
    df['kyles_lambda'] = calculate_kyles_lambda(df, window=lambda_window)
    
    # 确定预热期和正式数据的分界点
    warmup_end_idx = None
    if start_time_str and warmup_periods > 0:
        start_ts = pd.to_datetime(start_time_str)
        # 找到第一个大于等于开始时间的索引
        mask = df.index >= start_ts
        if mask.any():
            warmup_end_idx = df.index[mask][0]
    
    # 统计信息（只统计正式数据期，排除预热期）
    if warmup_end_idx is not None:
        valid_lambda = df.loc[warmup_end_idx:, 'kyles_lambda'].dropna()
        warmup_lambda = df.loc[:warmup_end_idx, 'kyles_lambda'].dropna()
    else:
        valid_lambda = df['kyles_lambda'].dropna()
        warmup_lambda = pd.Series(dtype=float)
    
    if len(valid_lambda) > 0:
        print(f"\nKyle's Lambda 统计 (正式数据期):")
        print(f"  平均值: {valid_lambda.mean():.6f}")
        print(f"  最大值: {valid_lambda.max():.6f}")
        print(f"  最小值: {valid_lambda.min():.6f}")
        print(f"  标准差: {valid_lambda.std():.6f}")
        print(f"  有效数据点: {len(valid_lambda)}")
    
    if len(warmup_lambda) > 0:
        print(f"\n预热期数据: {len(warmup_lambda)} 个数据点（Lambda 值不可用）")
    
    hover_texts = []
    for index, row in df.iterrows():
        lambda_val = row['kyles_lambda'] if pd.notna(row['kyles_lambda']) else 'N/A'
        is_warmup = (warmup_end_idx is not None and index < warmup_end_idx)
        warmup_status = " (预热期)" if is_warmup else ""
        text = (f"<b>Time</b>: {index.strftime('%Y-%m-%d %H:%M')}{warmup_status}<br>"
                f"<b>Open</b>: {row['open']:.2f}<br>"
                f"<b>High</b>: {row['high']:.2f}<br>"
                f"<b>Low</b>: {row['low']:.2f}<br>"
                f"<b>Close</b>: {row['close']:.2f}<br>"
                f"<b>Kyle's Lambda</b>: {lambda_val}")
        hover_texts.append(text)
    
    # 创建单个图表，使用双Y轴
    fig = go.Figure()
    
    # 1. 添加K线图（主Y轴，左侧）
    fig.add_trace(go.Candlestick(x=df.index,
                                 open=df['open'], high=df['high'],
                                 low=df['low'], close=df['close'],
                                 name='Price',
                                 yaxis='y',
                                 hoverinfo='text',
                                 hovertext=hover_texts))
    
    # 2. 添加 Kyle's Lambda（次Y轴，右侧）
    # 分别绘制正负值区域以便更好地区分
    lambda_positive = df['kyles_lambda'].copy()
    lambda_positive[lambda_positive < 0] = pd.NA
    lambda_negative = df['kyles_lambda'].copy()
    lambda_negative[lambda_negative >= 0] = pd.NA
    
    # 正值线（绿色）
    fig.add_trace(go.Scatter(x=df.index, y=lambda_positive, 
                             name="Kyle's Lambda (+)",
                             mode='lines+markers',
                             line=dict(color='green', width=2),
                             marker=dict(size=4, color='green'),
                             yaxis='y2',
                             hoverinfo='text',
                             hovertext=hover_texts,
                             showlegend=True))
    
    # 负值线（红色）
    fig.add_trace(go.Scatter(x=df.index, y=lambda_negative, 
                             name="Kyle's Lambda (-)",
                             mode='lines+markers',
                             line=dict(color='red', width=2),
                             marker=dict(size=4, color='red'),
                             yaxis='y2',
                             hoverinfo='text',
                             hovertext=hover_texts,
                             showlegend=True))
    
    # 添加零线（在次Y轴上）
    fig.add_hline(y=0, line_dash='dash', line_color='gray', 
                  opacity=0.5, yref='y2')
    
    # 添加预热期区域标注
    if warmup_end_idx is not None:
        # 在价格轴上添加预热期阴影
        fig.add_vrect(
            x0=df.index[0],
            x1=warmup_end_idx,
            fillcolor="yellow",
            opacity=0.15,
            layer="below",
            line_width=0,
            annotation_text="预热期",
            annotation_position="top left",
            annotation_font_size=10,
            annotation_font_color="orange"
        )
        
        # 添加分界线 - 使用 add_shape 替代 add_vline 以避免 Timestamp 类型问题
        fig.add_shape(
            type="line",
            x0=warmup_end_idx,
            x1=warmup_end_idx,
            y0=0,
            y1=1,
            yref="paper",
            line=dict(
                color="orange",
                width=2,
                dash="dash"
            ),
            opacity=0.7
        )
        
        # 添加标注文本
        fig.add_annotation(
            x=warmup_end_idx,
            y=1,
            yref="paper",
            text="预热期结束",
            showarrow=False,
            font=dict(size=10, color="orange"),
            xanchor="left",
            bgcolor="rgba(255,255,255,0.8)",
            bordercolor="orange",
            borderwidth=1
        )
    
    # 更新布局
    display_start = warmup_end_idx if warmup_end_idx is not None else df.index[0]
    start_str = display_start.strftime('%Y-%m-%d %H:%M')
    end_str = df.index[-1].strftime('%Y-%m-%d %H:%M')
    warmup_info = f" (预热期: {warmup_periods}根K线)" if warmup_periods > 0 else ""
    full_title = f'{chart_title} ({start_str} to {end_str}) with Kyle\'s Lambda (Window={lambda_window}){warmup_info}'
    
    fig.update_layout(
        title_text=full_title,
        xaxis=dict(
            title='Time',
            rangeslider=dict(visible=False)
        ),
        yaxis=dict(
            title='Price',
            side='left',
            domain=[0, 1]
        ),
        yaxis2=dict(
            title="Kyle's Lambda",
            side='right',
            overlaying='y',
            domain=[0, 1]
        ),
        hovermode='x unified',
        showlegend=True,
        legend=dict(
            x=1.02,
            y=1,
            xanchor='left',
            yanchor='top'
        )
    )
    
    fig.show()


if __name__ == '__main__':
    # ==================== 用户配置区 ====================
    BASE_DATA_PATH = "F:/personal/binance_klines"
    SYMBOL = "BTCUSDT"
    TIMEFRAME = "1h"
    START_TIME = "2025-06-05 00:00:00"
    END_TIME = "2025-08-07 00:00:00"
    LAMBDA_WINDOW = 48  # Kyle's Lambda 的滚动窗口大小
    # 预热期设置：自动加载 START_TIME 之前的数据用于指标计算
    # 预热期K线数量应该 >= LAMBDA_WINDOW，以确保指标从开始就有有效值
    WARMUP_PERIODS = LAMBDA_WINDOW  # 默认使用与窗口大小相同的预热期
    # ====================================================
    
    print("--- K线图表程序启动 (包含 Kyle's Lambda) ---")
    print(f"指标窗口大小: {LAMBDA_WINDOW}, 预热期: {WARMUP_PERIODS} 根K线")
    
    # 计算需要加载的数据时间范围（包含预热期）
    # 为了确保有足够的预热期数据，我们加载一个更早的时间范围
    start_dt = pd.to_datetime(START_TIME)
    # 根据时间周期计算预热期需要的时间范围（多加载一些以确保足够）
    timeframe_to_hours = {
        "1m": 1/60, "3m": 3/60, "5m": 5/60, "15m": 15/60, "30m": 30/60,
        "1h": 1, "2h": 2, "4h": 4, "6h": 6, "8h": 8, "12h": 12,
        "1d": 24, "3d": 72, "1w": 168
    }
    hours_per_period = timeframe_to_hours.get(TIMEFRAME, 1)
    # 多加载一些时间以确保有足够的数据（额外增加50%的缓冲）
    warmup_hours = int(WARMUP_PERIODS * hours_per_period * 1.5)
    warmup_delta = timedelta(hours=warmup_hours)
    
    data_start_time = start_dt - warmup_delta
    data_start_str = data_start_time.strftime('%Y-%m-%d %H:%M:%S')
    
    print(f"\n--- 正在加载数据 (包含预热期: {data_start_str} to {END_TIME}) ---")
    raw_data_df = load_data_in_range(data_start_str, END_TIME, BASE_DATA_PATH, SYMBOL, TIMEFRAME)
    final_df = process_and_filter_data(raw_data_df, START_TIME, END_TIME, 
                                       include_warmup=True, warmup_periods=WARMUP_PERIODS)
    
    if not final_df.empty:
        chart_title_str = f"{SYMBOL} - {TIMEFRAME}"
        plot_with_kyles_lambda(final_df, chart_title=chart_title_str, 
                              lambda_window=LAMBDA_WINDOW,
                              start_time_str=START_TIME,
                              warmup_periods=WARMUP_PERIODS)
        print("\n--- 图表生成完毕 ---")
    else:
        print("\n未能加载或处理任何图表数据，程序退出。请检查您的配置和文件路径。")

